#pragma once
#include "TFGenerator.h"
#include <ANN/ANN.h>
#include "MatlabMatrix.h"
#include "MatlabDrawer.h"
using namespace matlab; 

#include <boost/thread/mutex.hpp>
#include "MetricLearner.h"
#include "OriginalTFGenerator.h"

class FMLKRGeneralizer :
	public TFGenerator, MatlabDrawer
{
public:
	virtual ~FMLKRGeneralizer(void);

	FMLKRGeneralizer(MREAgent* p, int an, int od, taskspec_t& spec); 
	virtual void batchLearn( list<Transition>& history); 
	virtual void learn(const Transition* t); 
	virtual Observation predict(Observation st, Action a); 
	virtual Observation predict(Observation st, Action a, double& confidence); 
	virtual double getConfidence(Observation st, Action a);
	virtual void draw(); 


private:
	double predict(const Observation st, Action a, int dim, int k,  Observation_type& result, double& dists);	//predict the 'dim' dimension of the output
	void resetTrees();				//free any storage used by the learners
	
	void updatePoints(list<Transition>& history); 
	void updatePointsIndividual(std::list<Transition>& l, bool* use,  ANNpointArray& points, Observation_type* & targets,  int action, int dim); 
	
	void transformData();			//using the A matrix, transforms all the data to new coordinate system. 
	int transformData(ANNpointArray& data,Observation_type*& targets,  int size, int action, int dim); 
	
	void computeKernels(double* dists, double* kernels); 

	void drawScatteredPoints(int action, int startPos); 
	void drawOption1(); 
	void drawOption2(); 

public:
	int K;											//how many nearest neighbors to use 
	double epsilon;									//epsilon for approximate KNN
	double sigma2;									//the sigma^2 used to compute kernel
	int DRAW_OPTION; 

	MatlabMatrix<ANNkd_tree*> learners; 
	MatlabMatrix<ANNpointArray> dataPoints;	    	//local copy of the points (it should not be changed between execution of batchlearn). 
	MatlabMatrix<Observation_type*> targets;        //the labels for dataPoints
	MatlabMatrix<int> reducedDimensions; 
	int* dataSize;									//an int array of size action_number that specifies how many data point for each action is available
	ANNidx* nnIdx;									//this is a temporary place to put the indexes of the k nearest neighbors
	ANNdist* dists;									//this is a temporary place to put the distances of the k nn

	boost::mutex m_mutex; 

	//warning: just for debugging, remove
	OriginalTFGenerator* otf; 

	SingleVariateMetricLearner* metricLearner; 
	friend class MetricLearner; 
};
